import torch
import torch.nn as nn
import torch_geometric.nn as geom_nn
from .detectors.yolo_pose import YOLOv7Pose
from .trackers.deepsort_multimodal import DeepSORTMultimodal
from .st_gnn.st_gnn_model import STGNN
from .fusion_modules.cross_modal_fusion import CrossModalFusion

class NightPedestrianPredictor(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.detector = YOLOv7Pose(config['detector'])
        self.tracker = DeepSORTMultimodal(config['tracker'])
        self.fusion = CrossModalFusion(config['fusion'])
        self.st_gnn = STGNN(config['st_gnn'])
        self.intent_predictor = IntentPredictor(config['st_gnn']['hidden_dim'])
        self.trajectory_decoder = TrajectoryDecoder(config['st_gnn']['hidden_dim'])
        
    def forward(self, infrared_data, radar_data, context_data):
        """
        Args:
            infrared_data: [batch_size, seq_len, 1, H, W]
            radar_data: [batch_size, seq_len, num_points, 4]  # x, y, z, velocity
            context_data: dict with context features
        Returns:
            intent_probs: [batch_size, future_frames, 2]
            trajectories: [batch_size, future_frames, 2]
        """
        batch_size, seq_len = infrared_data.shape[:2]
        
        # 行人检测与跟踪
        detections = self.detector(infrared_data)  # [batch_size, seq_len, max_objects, 6+17*3]
        tracks = self.tracker(detections, radar_data)  # [batch_size, seq_len, num_tracks, track_dim]
        
        # 构建动态图结构
        graph_data = self.build_graph(tracks, context_data)
        
        # 多模态融合
        fused_features = self.fusion(graph_data.x, radar_data)
        
        # 时空图神经网络推理
        node_features = self.st_gnn(fused_features, graph_data.edge_index, graph_data.edge_attr)
        
        # 意图预测和轨迹解码
        intent_probs = self.intent_predictor(node_features)
        trajectories = self.trajectory_decoder(node_features)
        
        return intent_probs, trajectories, tracks
    
    def build_graph(self, tracks, context):
        """构建行人动态图结构"""
        # 实现图构建逻辑
        batch_size, seq_len, num_tracks, track_dim = tracks.shape
        
        # 提取节点特征
        node_features = []
        for b in range(batch_size):
            for t in range(seq_len):
                for i in range(num_tracks):
                    if tracks[b, t, i, 0] > 0:  # 有效的跟踪
                        feature = self.extract_node_features(tracks[b, t, i], context)
                        node_features.append(feature)
        
        # 构建边
        edge_index, edge_attr = self.build_edges(tracks, context)
        
        return GraphData(x=torch.stack(node_features), 
                        edge_index=edge_index, 
                        edge_attr=edge_attr)
    
    def extract_node_features(self, track, context):
        """提取节点特征"""
        # 运动特征
        pos = track[1:3]  # x, y
        vel = track[3:5]  # vx, vy
        acc = track[5:7]  # ax, ay
        
        # 姿态特征
        keypoints = track[7:7+51]  # 17个关键点 * 3
        pose_features = self.extract_pose_features(keypoints)
        
        # 行为特征
        behavior_features = self.extract_behavior_features(track, context)
        
        # 上下文特征
        context_features = self.extract_context_features(context)
        
        return torch.cat([pos, vel, acc, pose_features, behavior_features, context_features])
    
    def build_edges(self, tracks, context):
        """构建图边"""
        # 实现基于距离、运动一致性、视觉注意力的边构建
        pass

class IntentPredictor(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, output_dim=2):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, output_dim),
            nn.Softmax(dim=-1)
        )
    
    def forward(self, x):
        return self.network(x)

class TrajectoryDecoder(nn.Module):
    def __init__(self, input_dim, hidden_dim=128, output_dim=2):
        super().__init__()
        self.gru = nn.GRU(input_dim, hidden_dim, batch_first=True)
        self.linear = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        # x: [batch_size, num_nodes, features]
        batch_size, num_nodes, features = x.shape
        
        # 初始隐藏状态
        h0 = torch.zeros(1, batch_size * num_nodes, self.gru.hidden_size).to(x.device)
        
        # 重复特征用于多步预测
        x_repeated = x.unsqueeze(2).repeat(1, 1, self.config['training']['future_frames'], 1)
        x_flat = x_repeated.reshape(batch_size * num_nodes, 
                                  self.config['training']['future_frames'], 
                                  features)
        
        # GRU解码
        output, _ = self.gru(x_flat, h0)
        trajectories = self.linear(output)
        
        return trajectories.reshape(batch_size, num_nodes, 
                                  self.config['training']['future_frames'], 2)